import json
import re
import asyncio
import concurrent.futures
import time
from typing import List, Dict, Any, Callable, Optional

from agents.base_agent import BaseAgent
from my_utils.format_parser import remove_formulas_and_tables
from my_utils.llm_client import llm_chat, llm_chat_parallel
from my_utils.api_request import search_papers, query_by_title, query_by_title_contain
from my_utils.data_type import wrap_search_results
from prompts.paper_search import paper_search_user_prompt, filter_paper_user_prompt
from my_utils.data_type import QueryResult, display_query_results


def list_extract(raw_str: str):
    """
    用于从字符串中抽取出列表
    """
    # 使用正则匹配出列表字符串
    pattern = r'\[.*?\]'
    match = re.search(pattern, raw_str)
    # 获取匹配的列表字符串
    if not match:  # 新增检查：如果没有匹配到列表
        print(f"警告：在字符串中未能提取到列表格式: {raw_str}")
        return []  # 返回空列表或抛出异常，取决于需求
    list_str = match.group()
    # 将单引号替换为双引号,以便json加载
    list_str = list_str.replace("'", '"')
    # 使用json加载列表
    try:  # 新增检查：处理 JSON 加载错误
        result = json.loads(list_str)
        if not isinstance(result, list):  # 新增检查：确保结果是列表
            print(f"警告：提取的内容不是列表格式: {result}")
            return []
        return result
    except json.JSONDecodeError as e:
        print(f"错误：提取列表时JSON解析失败: {e}, 原始字符串: {list_str}")
        return []  # 返回空列表


class PaperSearchAgent(BaseAgent):
    def __init__(self, origin_query: str, user_summary_type: str, query_list: list = [], top_k=50,
                 papers_count: int = 50, distance_threshold=0.6, rerank=True):
        super().__init__()
        self.initial_query_list = query_list  # 保存初始列表以供异步初始化使用
        self.origin_query = origin_query
        self.query_history = []
        self.top_k = top_k
        self.distance_threshold = distance_threshold
        self.papers_count = papers_count
        self.user_summary_type = user_summary_type
        self.save_data = {}
        # origin_result 在 search_and_add 中填充
        self.origin_result = {}
        self.rerank = rerank

    # 异步初始化方法
    async def async_initialize(self):
        """ 异步执行初始化时的并行搜索 """
        self.query_history = list(set(self.initial_query_list + [self.origin_query]))  # 初始化查询历史，去重
        print(f"开始异步初始化搜索，查询列表: {self.query_history}")
        # 创建异步搜索任务
        tasks = [self.search_and_add(query, self.top_k) for query in self.query_history]
        # 使用 asyncio.gather 并发执行
        await asyncio.gather(*tasks)

    async def create_new_query(self):
        saved_papers_info = ""
        # index = 1 # index 未在循环外使用
        # 使用 enumerate 获取索引，避免手动管理 index
        for index, item in enumerate(self.origin_result.items()):
            # 假设 item[1] 是包含 'entity' 的字典
            if isinstance(item[1], dict) and 'entity' in item[1]:
                entity = item[1]['entity']
                paper_title = entity.get('paper_title', 'N/A')  # 使用 get 获取，提供默认值
                # chunk_text = entity.get('chunk_text', '') # chunk_text 未在此处使用
                saved_papers_info += f"已检索到的第{index + 1}篇论文块：\n标题: {paper_title}\n"
            else:
                print(f"警告：origin_result 中索引 {index} 的条目格式不正确: {item}")

        messages = [
            {"role": "user", "content": paper_search_user_prompt.format(user_origin_query=self.origin_query,
                                                                                 user_summary_type=self.user_summary_type,
                                                                                 query_history=json.dumps(
                                                                                     self.query_history,
                                                                                     ensure_ascii=False),
                                                                                 papers_info=saved_papers_info
                                                                                 )},
        ]
        llm_response = llm_chat(messages, temperature=0.6, model_name='GLM_4_PLUS')

        # 检查响应是否有效
        if not llm_response or not hasattr(llm_response, 'choices') or not llm_response.choices:
            print("错误：调用 create_new_query 的 llm_chat 未返回有效响应。")
            return []

        llm_result = llm_response.choices[0].message.content
        print(f"大模型编写的新查询:{llm_result}")
        try:
            result = list_extract(llm_result)  # 使用已有的 list_extract 函数
            unique_list = []
            # 确保 result 是列表
            if isinstance(result, list):
                for item in result:
                    # 添加检查，确保 item 是字符串或其他可哈希类型
                    if isinstance(item, (str, int, float, tuple)) and item not in self.query_history:
                        self.query_history.append(item)
                        unique_list.append(item)
                    elif item in self.query_history:
                        print(f"信息：查询 '{item}' 已存在于历史记录中，跳过。")
                    else:
                        print(f"警告：提取的新查询项类型不支持或不可哈希: {item} (类型: {type(item)})")

            return unique_list
        except Exception as e:
            print(f"从LLM响应中提取新查询列表时出错: {e}, 响应内容: {llm_result}")
            return []

    # parallel_filter_papers 保持异步，它调用了异步的 run_loop
    async def parallel_filter_papers(self, query, origin_search_result):
        filtered_papers = []
        # 检查 origin_search_result 是否为空或无效
        if not origin_search_result:
            print(f"警告：查询 '{query}' 的原始搜索结果为空，无法过滤。")
            return []

        # 准备用户输入列表
        user_inputs = []
        valid_indices = []  # 记录原始结果中有效条目的索引
        for idx, paper in enumerate(origin_search_result):
            # 增加检查，确保 paper 结构符合预期
            if isinstance(paper, dict) and 'entity' in paper and isinstance(paper['entity'], dict) and 'chunk_text' in \
                    paper['entity']:
                user_inputs.append(
                    filter_paper_user_prompt.format(query=query, chunk_text=paper['entity']['chunk_text']))
                valid_indices.append(idx)
            else:
                print(f"警告：查询 '{query}' 的原始搜索结果中，索引 {idx} 的条目格式不正确: {paper}")

        if not user_inputs:
            print(f"警告：查询 '{query}' 的原始搜索结果中没有找到有效的 chunk_text 进行过滤。")
            return []

        # 调用异步的 run_loop
        base_agent = BaseAgent()
        results = await base_agent.run_loop(
            model_name='GLM_4_PLUS',
            user_inputs=user_inputs,
            temperature=0
        )

        # 检查返回结果数量是否匹配
        if len(results) != len(user_inputs):
            print(
                f"错误：过滤查询 '{query}' 时，run_loop 返回的结果数量 ({len(results)}) 与输入数量 ({len(user_inputs)}) 不匹配。")
            return []  # 或者采取其他错误处理

        # 处理过滤结果
        for i, result in enumerate(results):
            original_index = valid_indices[i]  # 获取对应原始结果的索引
            pattern = r'<answer>\s*(.*?)\s*</answer>'
            # 确保 result 是字符串
            if not isinstance(result, str):
                print(f"警告：过滤查询 '{query}' 时，run_loop 返回了非字符串结果 (索引 {i}): {result}")
                continue  # 跳过非字符串结果

            match = re.search(pattern, result, re.DOTALL)
            if not match:
                # print(f"信息：过滤查询 '{query}' 的响应 (索引 {i}) 未匹配到 <answer> 标签: {result}")
                continue  # 如果没有明确的 True/False，跳过

            answer = match.group(1).strip().lower()  # 提取、去空白并转小写，便于比较
            # print(f"Debug: Query='{query}', Index={original_index}, Answer='{answer}'")

            if answer == 'true':
                filtered_papers.append(origin_search_result[original_index])
            # else: # 记录被过滤掉的论文
            #    print(f"信息：论文块 {original_index} (查询 '{query}') 被过滤，答案: {answer}")

        print(f"查询 '{query}' 过滤完成，保留 {len(filtered_papers)}/{len(valid_indices)} 个有效论文块。")
        return filtered_papers

    async def targeted_paper_collection(self):
        loop_count = 0  # 防止潜在的无限循环
        max_loops = 10  # 设定一个最大循环次数，避免因无法达到 papers_count 而死循环

        while len(self.origin_result) < self.papers_count and loop_count < max_loops:
            loop_count += 1
            print(f"收集第 {loop_count} 轮，当前已收集 {len(self.origin_result)}/{self.papers_count} 篇。")
            start_time = time.time()

            start_time_create_new_query = time.time()
            # 调用异步的 create_new_query
            new_query_list = await self.create_new_query()
            end_time_create_new_query = time.time()
            print(f"第 {loop_count} 轮生成新查询用时：", int((end_time_create_new_query - start_time_create_new_query) // 60), "分钟",
                  int((end_time_create_new_query - start_time_create_new_query) % 60), "秒")

            if not new_query_list:
                print("未能生成新的查询，定向收集可能提前结束。")
                break  # 如果无法生成新查询，则跳出循环

            # 为新查询创建异步搜索任务
            tasks = [self.search_and_add(final_query) for final_query in new_query_list]
            # 使用 asyncio.gather 并发执行搜索和添加
            # 添加错误处理，防止一个任务失败导致 gather 整体失败
            results = await asyncio.gather(*tasks, return_exceptions=True)

            # 检查 gather 的结果
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    print(f"执行查询 '{new_query_list[i]}' 的 search_and_add 任务时出错: {result}")
                    # 可以选择记录错误，但循环继续
            end_time = time.time()
            print(f"第 {loop_count} 轮检索用时：", int((end_time - start_time) // 60), "分钟", int((end_time - start_time) % 60), "秒")

            # 检查是否达到数量要求，如果在一轮结束后就达到，可以提前退出
            if len(self.origin_result) >= self.papers_count:
                print(f"已收集到 {len(self.origin_result)} 篇，达到目标 {self.papers_count}，停止收集。")
                break

        if loop_count >= max_loops:
            print(f"警告：定向收集达到最大循环次数 {max_loops}，当前收集 {len(self.origin_result)} 篇。")

        try:
            # 确保目录存在等操作可以在这里添加
            file_path = 'examples_data/origin_paper_search_result.json'
            print(f"准备将 {len(self.origin_result)} 条结果写入 {file_path}...")
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(self.origin_result, f, ensure_ascii=False, indent=4)
            print(f"结果成功写入 {file_path}")
        except IOError as e:
            print(f"错误：写入结果文件失败: {e}")
        except TypeError as e:
            print(f"错误：结果数据无法序列化为 JSON: {e}")

        return self.origin_result

    async def search_and_add(self, query, top_k=40):
        try:
            print(f"开始为查询 '{query}' 搜索和添加论文...")

            # Python 3.9以上
            origin_search_result = await asyncio.to_thread(search_papers, query, top_k)

            print(f"查询 '{query}' 原始搜索到 {len(origin_search_result)} 个结果。")

            # 对搜索结果进行处理（remove_formulas_and_tables），并保存回origin_search_result
            # origin_search_result = [remove_formulas_and_tables(item) for item in origin_search_result]

            for idx, item in enumerate(origin_search_result):
                try:
                    # 检查结构
                    if isinstance(item, dict) and 'entity' in item and isinstance(
                            item['entity'], dict) and 'chunk_text' in item['entity']:
                        original_text = item['entity']['chunk_text']
                        cleaned_text = remove_formulas_and_tables(original_text)
                        item['entity']['chunk_text'] = ' '.join(cleaned_text.split()[:1800])
                    else:
                        print(f"警告：结果条目 {idx} 格式不正确，无法进行后处理。")
                except Exception as e:
                    print(f"错误：处理条目 {idx} 时移除公式/表格失败: {e}")

            # 调用异步的过滤函数
            filtered_papers = await self.parallel_filter_papers(query, origin_search_result)

            added_count = 0  # 记录本次添加的数量
            # 对过滤后的结果进行处理
            for paper in filtered_papers:
                # 再次检查 paper 结构
                if isinstance(paper, dict) and 'entity' in paper and isinstance(paper['entity'], dict) \
                        and 'paper_id' in paper['entity'] and 'chunk_id' in paper['entity']:

                    chunk_index = str(paper['entity']['paper_id']) + "_" + str(paper['entity']['chunk_id'])

                    if chunk_index not in self.origin_result:
                        self.origin_result[chunk_index] = paper
                        added_count += 1
                else:
                    print(f"警告：过滤后的论文条目格式不正确，无法添加到结果集: {paper}")

            print(f"查询 '{query}' 处理完成，新增 {added_count} 个论文块到结果集。当前总数: {len(self.origin_result)}")

        except Exception as e:
            print(f"处理查询 '{query}' 的 search_and_add 时发生错误: {e}")
            # 重新抛出异常


async def search_paper_loop(origin_query: str, user_summary_type: str, keywords: list = [], papers_count: int = 50,
                            top_k=40, rerank=False):
    # 实例化 Agent
    agent = PaperSearchAgent(origin_query=origin_query, user_summary_type=user_summary_type, query_list=keywords,
                             top_k=top_k, papers_count=papers_count, rerank=rerank)

    print("文献收集中...")
    start_time_1 = time.time()
    # 调用异步初始化方法
    await agent.async_initialize()

    end_time_1 = time.time()
    print("第一轮检索用时：", int((end_time_1 - start_time_1) // 60), "分钟", int((end_time_1 - start_time_1) % 60), "秒")

    # 调用异步的定向收集方法
    origin_result = await agent.targeted_paper_collection()

    print("文献收集完成。")
    print("搜索过的问题:", agent.query_history)

    return origin_result


if __name__ == "__main__":
    origin_query = "What is the current status and challenges of Text2SQL research?"
    user_summary_type = "A review of the current research status in a certain direction"
    papers_count = 100
    origin_result = search_paper_loop(origin_query,user_summary_type ,papers_count=papers_count,top_k=30,rerank=True)